import argparse
import numpy as np
from deepsignal3.utils.process_utils import base2code_dna
from deepsignal3.utils.process_utils import get_motif_seqs

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/rice_ribenqing/test/20x_1/pod5.CHH.CG.CHG.features.tsv",
        help="trainset input address.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/rice_ribenqing/test/20x_1/pod5.CHH.features.tsv",
        help="trainset output address.",
    )
    parser.add_argument(
        "--motif",
        type=str,
        default="CHH",
        help="motif",
    )
    return parser.parse_args()



def parse_data(args,motif_seq):
    fh = open(args.output, "w", buffering=512)
    with open(args.input, "r") as f:
        for line in f:
            words = line.strip().split()
            motifset = set(motif_seq)
            motiflen = len(list(motifset)[0])
            if words[6][len(words[6])// 2:len(words[6])// 2+ motiflen] not in motifset:
                continue            
            fh.write(line)
            
            

if __name__ == '__main__':
    args=parse_args()
    is_dna=True
    motif_seq = get_motif_seqs(args.motif, is_dna)
    parse_data(args,motif_seq)